#####
# Script to generate figure 2.
#####

library(ggplot2)
library(cowplot)
library(patchwork)
library(viridis)
library(tidyverse)
library(pracma)
library(nleqslv)
library(rootSolve)
library(fields)
library(doParallel)
library(scales)
library(data.table)

source("equations.r")
source("simulation_functions.r")
source("planner_simulation.r")

avg_sat_decay <- 1 # 1: satellites are infinitely lived
p <- 1
F <- 35
discount <- 0.99
r <- ((1 - discount)/(discount))
fe_eqm <- p/F - r #the equilibrium risk
d <- 0.5
m <- 0
aSS <- 0.005
aSD <- 0.005
aDD <- 0.025 
bSS <- 10
bSD <- 5.375
bDD <- 0.9 # 0.075
T <- 150
label <- "benchmark"

## Calibration scale factors.
s_scale_factor <- 68289.84/3.728625
d_scale_factor <- 367372.3/4.214421

params <- data.frame(avg_sat_decay=avg_sat_decay, p=p, F=F, discount=discount, r=r, fe_eqm=fe_eqm, d=d, m=m, aSS=aSS, aSD=aSD, aDD=aDD, bSS=bSS, bSD=bSD, bDD=bDD, T=T)
params2 <- params
params2$fe_eqm <- params2$fe_eqm*2

# generate a satellite-debris grid
S_element <- (seq(0,10,length.out=300))
D_element <- (seq(0,10,length.out=300))
SD_grid <- expand.grid(S_element,D_element)
colnames(SD_grid) <- c("satellites","debris")

# calculate equilibrium values
fig1_dfrm <- data.frame(SD_grid) %>% 
				mutate(oa_launches = eqm_launch_rate(satellites,debris,params)) %>%
				mutate(oa_S_ = S_(oa_launches,satellites,debris,params)) %>%
				mutate(oa_D_ = D_(oa_launches,satellites,debris,params)) %>%
				mutate(numerical_coll_rate = L(oa_S_,oa_D_,params)) %>%
				mutate(oa_coll_rate = eqm_sat_stock(debris,params)) %>%
				mutate(oa_coll_rate_2 = eqm_sat_stock(debris,params2))

head(fig1_dfrm)

constrained_launch_path <- OA.ts(0,0,150,0,opt=0,constraint=0,params=params2)
clp_sats <- as.data.frame(constrained_launch_path)$satellites
clp_debs <- as.data.frame(constrained_launch_path)$debris
length(clp_sats) <- nrow(fig1_dfrm)
length(clp_debs) <- nrow(fig1_dfrm)

unconstrained_launch_path <- OA.ts(0,0,150,0,opt=0,constraint=0,params=params)
ulp_sats <- as.data.frame(unconstrained_launch_path)$satellites
ulp_debs <- as.data.frame(unconstrained_launch_path)$debris
length(ulp_sats) <- nrow(fig1_dfrm)
length(ulp_debs) <- nrow(fig1_dfrm)

fig1_dfrm <- cbind(fig1_dfrm,clp_sats,clp_debs)

fig1_time_dfrm <- left_join( as.data.frame(constrained_launch_path), as.data.frame(unconstrained_launch_path), by=c("time"), suffix=c("_clp","_ulp")) %>%
	mutate(
		satellites_clp = satellites_clp*s_scale_factor,
		satellites_ulp = satellites_ulp*s_scale_factor,
		launches_clp = launches_clp*s_scale_factor,
		launches_ulp = launches_ulp*s_scale_factor,
		debris_clp = debris_clp*d_scale_factor,
		debris_ulp = debris_ulp*d_scale_factor
	)

# brewer.pal(n = 8, name = "Accent")

figure_1time_a <- ggplot(data=fig1_time_dfrm[1:25,], aes(x=time)) +
				geom_line(aes(y=launches_clp), linetype="dashed", color="grey39", size=1) + 
				geom_line(aes(y=launches_ulp), linetype="dotted", color="grey39", size=1) + 
				geom_segment(x=0,y=0, xend=35,yend=0, size=0.15) + geom_segment(x=0,y=0, xend=0,yend=5, size=0.15) +
				ylab("Launches") + xlab("Time") +
				scale_x_continuous(labels=comma) + 
				scale_y_continuous(labels=comma) +
				theme_bw() + theme(text=element_text(family="Arial", size=15))

figure_1time_b <- ggplot(data=fig1_time_dfrm[1:25,], aes(x=time)) +
				geom_line(aes(y=risk_clp), linetype="dashed", color="grey39", size=1.5) + 
				geom_line(aes(y=risk_ulp), linetype="dotted", color="grey39", size=1.5) +  
				geom_hline(aes(yintercept=params2$fe_eqm), linetype="dashed", color="#7FC97F", size=1) +
				geom_hline(aes(yintercept=params$fe_eqm), linetype="dotted", color="#7FC97F", size=1) +
				ylab("Collision risk") + xlab("Time") +
				geom_segment(x=0,y=0, xend=35,yend=0, size=0.15) + geom_segment(x=0,y=0, xend=0,yend=5, size=0.15) +
				scale_x_continuous(labels=comma) + 
				scale_y_continuous(labels=comma) +
				theme_bw() + theme(text=element_text(family="Arial", size=15))

figure_1time_c <- ggplot(data=fig1_time_dfrm[1:25,], aes(x=time)) +
				geom_line(aes(y=satellites_clp), linetype="dashed", color="grey39", size=1) + 
				geom_line(aes(y=satellites_ulp), linetype="dotted", color="grey39", size=1) +  
				ylab("Satellites") + xlab("Time") +
				geom_segment(x=0,y=0, xend=35,yend=0, size=0.15) + geom_segment(x=0,y=0, xend=0,yend=5, size=0.15) +
				scale_x_continuous(labels=comma) + 
				scale_y_continuous(labels=comma) +
				theme_bw() + theme(text=element_text(family="Arial", size=15))


figure_1time_d <- ggplot(data=fig1_time_dfrm[1:25,], aes(x=time)) +
				geom_line(aes(y=debris_clp), linetype="dashed", color="grey39", size=1) + 
				geom_line(aes(y=debris_ulp), linetype="dotted", color="grey39", size=1) +  
				ylab("Debris") + xlab("Time") +
				geom_segment(x=0,y=0, xend=35,yend=0, size=0.15) + geom_segment(x=0,y=0, xend=0,yend=5, size=0.15) +
				scale_x_continuous(labels=comma) + 
				scale_y_continuous(labels=comma) +
				theme_bw() + theme(text=element_text(family="Arial", size=15))


fig1_time_panel <- plot_grid(figure_1time_a, figure_1time_b, figure_1time_c, figure_1time_d, align="h", labels=c("a","b","c","d"), nrow=2)

ggsave(
	paste0("../images/figure-2--example-openaccess-trajectories.png"),
	((figure_1time_a / figure_1time_c) | (figure_1time_b / figure_1time_d)) + plot_annotation(tag_levels = 'A') & plot_layout(guides = "collect") & theme(legend.position = 'right'),
	width = 16*2/3,
	height = 9*2/3,
	units = "in",
	dpi = 300)

